﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Runtime.InteropServices;
using System.Net;
using System.Diagnostics;

namespace iMan
{
    public class IpHelper
    {
        private const int NO_ERROR = 0;
        private const int TCP_STATE_CLOSED = 1;
        private const int TCP_STATE_LISTEN = 2;
        private const int TCP_STATE_SYN_SENT = 3;
        private const int TCP_STATE_SYN_RCVD = 4;
        private const int TCP_STATE_ESTAB = 5;
        private const int TCP_STATE_FIN_WAIT1 = 6;
        private const int TCP_STATE_FIN_WAIT2 = 7;
        private const int TCP_STATE_CLOSE_WAIT = 8;
        private const int TCP_STATE_CLOSING = 9;
        private const int TCP_STATE_LAST_ACK = 10;
        private const int TCP_STATE_TIME_WAIT = 11;
        private const int TCP_STATE_DELETE_TCB = 12;

        protected MIB_TCPTABLE tcpTable;
        protected MIB_TCPSTATS tcpStats;
        protected MIB_EXTCPTABLE tcpExTable;
        protected MIB_TCPTABLE_OWNER_PID tcpExtendedTable;
        protected MIB_UDPTABLE udpTable;
        protected MIB_UDPSTATS udpStats;
        protected MIB_EXUDPTABLE udpExTable;
        protected MIB_UDPTABLE_OWNER_PID udpExtendedTable;

        public IpHelper() { }

        public MIB_TCPTABLE TcpTable
        {
            get { return tcpTable; }
        }

        public MIB_TCPSTATS TcpStats
        {
            get { return tcpStats; }
        }

        public MIB_EXTCPTABLE TcpExTable
        {
            get { return tcpExTable; }
        }

        public MIB_TCPTABLE_OWNER_PID TcpExtendedTable
        {
            get { return tcpExtendedTable; }
        }

        public MIB_UDPTABLE UdpTable
        {
            get { return udpTable; }
        }

        public MIB_UDPSTATS UdpStats
        {
            get { return udpStats; }
        }

        public MIB_EXUDPTABLE UdpExTable
        {
            get { return udpExTable; }
        }

        public MIB_UDPTABLE_OWNER_PID UdpExtendedTable
        {
            get { return udpExtendedTable; }
        }

        public void GetTcpStats()
        {
            tcpStats = new MIB_TCPSTATS();
            IpHlpAPIWrapper.GetTcpStatistics(ref tcpStats);
        }

        public void GetUdpStats()
        {
            udpStats = new MIB_UDPSTATS();
            IpHlpAPIWrapper.GetUdpStatistics(ref udpStats);
        }

        public void GetExTcpTable()
        {
            int rowSize = 24;
            int bufferSize = 100000;
            IntPtr dummyMem = Marshal.AllocHGlobal(bufferSize);
            int res = IpHlpAPIWrapper.AllocateAndGetTcpExTableFromStack(ref dummyMem, true, IpHlpAPIWrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                throw new InvalidOperationException("Error get TcpExTable from stack");
            }
            int entries = (int)Marshal.ReadIntPtr(dummyMem);
            dummyMem = IntPtr.Zero;
            Marshal.FreeHGlobal(dummyMem);

            bufferSize = entries * rowSize + 4;
            tcpExTable = new MIB_EXTCPTABLE();
            IntPtr lpTable = Marshal.AllocHGlobal(bufferSize);
            res = IpHlpAPIWrapper.AllocateAndGetTcpExTableFromStack(ref lpTable, true, IpHlpAPIWrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                throw new InvalidOperationException("Error get TcpExTable from stack");
            }
            IntPtr current = lpTable;
            int currentIndex = 0;
            entries = (int)Marshal.ReadIntPtr(current);
            tcpExTable.table = new MIB_EXTCPROW[entries];
            currentIndex += 4;
            current = (IntPtr)((int)current + currentIndex);
            for (int i = 0; i < entries; i++)
            {
                tcpExTable.table[i].StrgState = ConvertState((int)Marshal.ReadIntPtr(current));
                tcpExTable.table[i].iState = (int)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                UInt32 localAddr = (UInt32)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                UInt32 localPort = (UInt32)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                tcpExTable.table[i].Local = new IPEndPoint(localAddr, (int)ConvertPort(localPort));
                UInt32 remoteAddr = (UInt32)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                UInt32 remotePort = 0;
                if (remoteAddr != 0)
                {
                    remotePort = (UInt32)Marshal.ReadIntPtr(current);
                    remotePort = ConvertPort(remotePort);
                }
                tcpExTable.table[i].Remote = new IPEndPoint(remoteAddr, (int)remotePort);
                current = (IntPtr)((int)current + 4);
                tcpExTable.table[i].dwProcessId = (int)Marshal.ReadIntPtr(current);
                tcpExTable.table[i].ProcessName = GetProcessName(tcpExTable.table[i].dwProcessId);
                current = (IntPtr)((int)current + 4);
            }

            Marshal.FreeHGlobal(lpTable);
            current = IntPtr.Zero;
        }

        public void GetExUdpTable()
        {
            int rowSize = 12;
            int bufferSize = 100000;
            IntPtr lpTable = Marshal.AllocHGlobal(bufferSize);
            int res = IpHlpAPIWrapper.AllocateAndGetUdpExTableFromStack(ref lpTable, true, IpHlpAPIWrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                throw new InvalidOperationException("Error get udp table");
            }
            int entries = (int)Marshal.ReadIntPtr(lpTable);
            lpTable = IntPtr.Zero;
            Marshal.FreeHGlobal(lpTable);

            bufferSize = entries * rowSize + 4;
            udpExTable = new MIB_EXUDPTABLE();
            lpTable = Marshal.AllocHGlobal(bufferSize);
            res = IpHlpAPIWrapper.AllocateAndGetUdpExTableFromStack(ref lpTable, true, IpHlpAPIWrapper.GetProcessHeap(), 0, 2);
            if (res != NO_ERROR)
            {
                throw new InvalidOperationException("Error get udp table");
            }
            IntPtr current = lpTable;
            int CurrentIndex = 0;
            entries = (int)Marshal.ReadIntPtr(current);
            udpExTable.dwNumEntries = entries;
            udpExTable.table = new MIB_EXUDPROW[entries];
            CurrentIndex += 4;
            current = (IntPtr)((int)current + CurrentIndex);
            for (int i = 0; i < entries; i++)
            {
                UInt32 localAddr = (UInt32)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                UInt32 localPort = (UInt32)Marshal.ReadIntPtr(current);
                current = (IntPtr)((int)current + 4);
                udpExTable.table[i].Local = new IPEndPoint(localAddr, ConvertPort(localPort));
                udpExTable.table[i].dwProcessId = (int)Marshal.ReadIntPtr(current);
                udpExTable.table[i].ProcessName = GetProcessName(udpExTable.table[i].dwProcessId);
                current = (IntPtr)((int)current + 4);
            }
            Marshal.FreeHGlobal(lpTable);
            current = IntPtr.Zero;
        }

        public void GetTcpTable()
        {
            byte[] buffer = new byte[20000];
            int pdwSize = 20000;
            int res = IpHlpAPIWrapper.GetTcpTable(buffer, out pdwSize, true);
            if (res != NO_ERROR)
            {
                buffer = new byte[pdwSize];
                res = IpHlpAPIWrapper.GetTcpTable(buffer, out pdwSize, true);
                if (res != 0) throw new InvalidOperationException("Error get tcp table");
            }

            tcpTable = new MIB_TCPTABLE();
            int nOffset = 0;
            tcpTable.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            nOffset += 4;
            tcpTable.table = new MIB_TCPROW[tcpTable.dwNumEntries];

            for (int i = 0; i < tcpTable.dwNumEntries; i++)
            {
                int st = Convert.ToInt32(buffer[nOffset]);
                tcpTable.table[i].StrgState = ConvertState(st);
                tcpTable.table[i].iState = st;
                nOffset += 4;
                string LocalAdrr = buffer[nOffset].ToString() + "." + buffer[nOffset + 1].ToString() + "." + buffer[nOffset + 2].ToString() + "." + buffer[nOffset + 3].ToString();
                nOffset += 4;
                int LocalPort = (((int)buffer[nOffset]) << 8) + (((int)buffer[nOffset + 1])) +
                    (((int)buffer[nOffset + 2]) << 24) + (((int)buffer[nOffset + 3]) << 16);

                nOffset += 4;
                tcpTable.table[i].Local = new IPEndPoint(IPAddress.Parse(LocalAdrr), LocalPort);
                string RemoteAdrr = buffer[nOffset].ToString() + "." + buffer[nOffset + 1].ToString() + "." + buffer[nOffset + 2].ToString() + "." + buffer[nOffset + 3].ToString();
                nOffset += 4;
                int RemotePort;
                if (RemoteAdrr == "0.0.0.0")
                {
                    RemotePort = 0;
                }
                else
                {
                    RemotePort = (((int)buffer[nOffset]) << 8) + (((int)buffer[nOffset + 1])) +
                        (((int)buffer[nOffset + 2]) << 24) + (((int)buffer[nOffset + 3]) << 16);
                }
                nOffset += 4;
                tcpTable.table[i].Remote = new IPEndPoint(IPAddress.Parse(RemoteAdrr), RemotePort);
            }
        }

        public void GetExtendedTcpTable()
        {
            int AF_INET = 2;
            int buffSize = 20000;
            byte[] buffer = new byte[buffSize];

            int res = IpHlpAPIWrapper.GetExtendedTcpTable(buffer, out buffSize, true, AF_INET, TCP_TABLE_CLASS.TCP_TABLE_OWNER_PID_ALL, 0);
            if (res != NO_ERROR) throw new InvalidOperationException("Error get extended tcp table");

            tcpExtendedTable = new MIB_TCPTABLE_OWNER_PID();
            int nOffset = 0;
            tcpExtendedTable.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            nOffset += 4;
            tcpExtendedTable.table = new MIB_TCPROW_OWNER_PID[tcpExtendedTable.dwNumEntries];
            for (int i = 0; i < tcpExtendedTable.dwNumEntries; i++)
            {
                int st = Convert.ToInt32(buffer[nOffset]);
                tcpExtendedTable.table[i].State = ConvertState(st);
                tcpExtendedTable.table[i].dwState = st;
                nOffset += 4;
                tcpExtendedTable.table[i].Local = BufferToIPEndPoint(buffer, ref nOffset, false);
                tcpExtendedTable.table[i].Remote = BufferToIPEndPoint(buffer, ref nOffset, true);
                tcpExtendedTable.table[i].dwOwningPid = BufferToInt(buffer, ref nOffset);
                tcpExtendedTable.table[i].ProcessName = GetProcessName(tcpExtendedTable.table[i].dwOwningPid);
            }
        }

        public void GetUdpTable()
        {
            byte[] buffer = new byte[20000];
            int pdwSize = 20000;
            int res = IpHlpAPIWrapper.GetUdpTable(buffer, out pdwSize, true);
            if (res != NO_ERROR)
            {
                buffer = new byte[pdwSize];
                res = IpHlpAPIWrapper.GetUdpTable(buffer, out pdwSize, true);
                if (res != 0) throw new InvalidOperationException("Error get udp table");
            }

            udpTable = new MIB_UDPTABLE();
            int nOffset = 0;
            udpTable.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            nOffset += 4;
            udpTable.table = new MIB_UDPROW[udpTable.dwNumEntries];
            for (int i = 0; i < udpTable.dwNumEntries; i++)
            {
                string LocalAdrr = buffer[nOffset].ToString() + "." + buffer[nOffset + 1].ToString() + "." + buffer[nOffset + 2].ToString() + "." + buffer[nOffset + 3].ToString();
                nOffset += 4;

                int LocalPort = (((int)buffer[nOffset]) << 8) + (((int)buffer[nOffset + 1])) +
                    (((int)buffer[nOffset + 2]) << 24) + (((int)buffer[nOffset + 3]) << 16);
                nOffset += 4;
                udpTable.table[i].Local = new IPEndPoint(IPAddress.Parse(LocalAdrr), LocalPort);
            }
        }

        public void GetExtendedUdpTable()
        {
            int AF_INET = 2;
            int buffSize = 20000;
            byte[] buffer = new byte[buffSize];

            int res = IpHlpAPIWrapper.GetExtendedUdpTable(buffer, out buffSize, true, AF_INET, UDP_TABLE_CLASS.UDP_TABLE_OWNER_PID, 0);
            if (res != NO_ERROR) throw new InvalidOperationException("Error get extended udp table");

            udpExtendedTable = new MIB_UDPTABLE_OWNER_PID();
            int nOffset = 0;
            udpExtendedTable.dwNumEntries = Convert.ToInt32(buffer[nOffset]);
            nOffset += 4;
            udpExtendedTable.table = new MIB_UDPROW_OWNER_PID[udpExtendedTable.dwNumEntries];
            for (int i = 0; i < udpExtendedTable.dwNumEntries; i++)
            {
                udpExtendedTable.table[i].Local = BufferToIPEndPoint(buffer, ref nOffset, false);
                udpExtendedTable.table[i].dwOwningPid = BufferToInt(buffer, ref nOffset);
                udpExtendedTable.table[i].ProcessName = GetProcessName(udpExtendedTable.table[i].dwOwningPid);
            }
        }

        private UInt16 ConvertPort(UInt32 dwPort)
        {
            byte[] b = new byte[2];
            b[0] = byte.Parse((dwPort >> 8).ToString());
            b[1] = byte.Parse((dwPort & 0xFF).ToString());
            return BitConverter.ToUInt16(b, 0);
        }

        private string ConvertState(int state)
        {
            string res = "";
            switch (state)
            {
                case TCP_STATE_CLOSED: res = "CLOSED"; break;
                case TCP_STATE_LISTEN: res = "LISTEN"; break;
                case TCP_STATE_SYN_SENT: res = "SYN_SENT"; break;
                case TCP_STATE_SYN_RCVD: res = "SYN_RCVD"; break;
                case TCP_STATE_ESTAB: res = "ESTAB"; break;
                case TCP_STATE_FIN_WAIT1: res = "FIN_WAIT1"; break;
                case TCP_STATE_FIN_WAIT2: res = "FIN_WAIT2"; break;
                case TCP_STATE_CLOSE_WAIT: res = "CLOSE_WAIT"; break;
                case TCP_STATE_CLOSING: res = "CLOSING"; break;
                case TCP_STATE_LAST_ACK: res = "LAST_ACK"; break;
                case TCP_STATE_TIME_WAIT: res = "TIME_WAIT"; break;
                case TCP_STATE_DELETE_TCB: res = "DELETE_TCB"; break;
            }
            return res;
        }

        private IPEndPoint BufferToIPEndPoint(byte[] buffer, ref int nOffset, bool IsRemote)
        {
            Int64 m_Address = ((((buffer[nOffset + 3] << 0x18) | (buffer[nOffset + 2] << 0x10)) | (buffer[nOffset + 1] << 8)) | buffer[nOffset]) & ((long)0xffffffff);
            nOffset += 4;
            int m_Port = 0;
            m_Port = (IsRemote && (m_Address == 0)) ? 0 :
                        (((int)buffer[nOffset]) << 8) + (((int)buffer[nOffset + 1])) + (((int)buffer[nOffset + 2]) << 24) + (((int)buffer[nOffset + 3]) << 16);
            nOffset += 4;

            IPEndPoint temp = new IPEndPoint(m_Address, m_Port);
            return temp;
        }

        private int BufferToInt(byte[] buffer, ref int nOffset)
        {
            int res = (((int)buffer[nOffset])) + (((int)buffer[nOffset + 1]) << 8) +
                (((int)buffer[nOffset + 2]) << 16) + (((int)buffer[nOffset + 3]) << 24);
            nOffset += 4;
            return res;
        }

        private string GetProcessName(int pid)
        {
            try
            {
                Process p = Process.GetProcessById(pid);
                return p.ProcessName;
            }
            catch
            {
                return "Unknown";
            }
        }
    }
}
